// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <gtest/gtest.h>
#include <tests/vector/vector_value_generator.h>
#include <pollux/vector/vector_maker.h>
#include <pollux/vector/vector_maker_stats.h>
#include <pollux/testing/vector/vector_test_base.h>

namespace kumo::pollux::test {
    inline std::vector<VectorEncoding::Simple> kFullValueTypes = {
        VectorEncoding::Simple::DICTIONARY,
        VectorEncoding::Simple::FLAT,
        VectorEncoding::Simple::SEQUENCE,
        VectorEncoding::Simple::CONSTANT
    };

    inline std::vector<VectorEncoding::Simple> kNonConstantTypes = {
        VectorEncoding::Simple::DICTIONARY,
        VectorEncoding::Simple::FLAT,
        VectorEncoding::Simple::SEQUENCE,
        VectorEncoding::Simple::BIASED
    };

    inline std::vector<VectorEncoding::Simple> kAllTypes = {
        VectorEncoding::Simple::DICTIONARY,
        VectorEncoding::Simple::FLAT,
        VectorEncoding::Simple::SEQUENCE,
        VectorEncoding::Simple::BIASED,
        VectorEncoding::Simple::CONSTANT
    };

    template<typename T>
    using ExpectedData = std::vector<std::optional<T> >;

    template<typename T>
    class VectorGeneratedData {
    public:
        void sort() {
            std::sort(data_.begin(), data_.end());
        }

        void push_back(std::optional<T> value) {
            if (value.has_value()) {
                auto ownedValue = stringViewBufferHolder_.getOwnedValue(value.value());
                data_.push_back(std::move(ownedValue));
            } else {
                data_.push_back(std::nullopt);
            }
        }

        void pop_back() {
            data_.pop_back();
        }

        void reserve(size_t amount) {
            data_.reserve(amount);
        }

        const ExpectedData<T> &data() const {
            return data_;
        }

        StringViewBufferHolder &stringViewBufferHolder() {
            return stringViewBufferHolder_;
        }

    protected:
        ExpectedData<T> data_;

        // In case T is StringView, the buffer below holds their actual data.
        std::shared_ptr<memory::MemoryPool> scopedPool =
                memory::memoryManager()->addLeafPool();
        StringViewBufferHolder stringViewBufferHolder_ =
                StringViewBufferHolder(scopedPool.get());
    };

    template<typename T>
    std::optional<T> getVal(
        const ExpectedData<T> &data,
        int32_t count,
        int32_t idx,
        int32_t sequenceCount,
        int32_t sequenceLength) {
        bool useSequences = sequenceCount > 0 && sequenceLength > 1;
        int64_t idxToUse = idx;
        if (useSequences) {
            auto chunkWidth = (count / data.size());
            idxToUse = idx / chunkWidth;
        }

        return data[idxToUse % data.size()];
    }

    /// Pass a fixed number when we want to generate determined pseudo-random
    /// data. In order to generate different sets of determined pseudo-random
    /// data, we can pass different integers at callsites. e.g. magic numbers
    /// 1000 for callsite A, 1070 for callsite B and etc.
    /// Otherwise generate unpredictable pseudo-random for benchmark and perf
    /// when seed is std::nullopt.
    ///
    /// `fixedWidthStringSize`: std::nullopt means generating variable-width string
    /// elements. A number means generating fixed-width string elements.
    template<typename T>
    VectorGeneratedData<T> genTestData(
        int32_t cardinality,
        const TypePtr type = CppToType<T>::create(),
        bool includeNulls = false,
        bool sorted = false,
        bool useFullTypeRange = false,
        const std::optional<uint32_t> &seed = std::nullopt,
        const std::optional<uint32_t> &fixedWidthStringSize = std::nullopt) {
        // Generate unpredictable random data by default.
        std::optional<melon::Random::DefaultGenerator> rng = std::nullopt;

        if (seed.has_value()) {
            rng = melon::Random::create();
            // Use the given seed so that it can generate predictable random data.
            rng.value().seed(seed.value());
        }

        VectorGeneratedData<T> testData;
        testData.reserve(cardinality);

        for (int64_t m = 0; m < cardinality; m++) {
            testData.push_back(VectorValueGenerator::cardValueOf<T>(
                useFullTypeRange,
                rng,
                testData.stringViewBufferHolder(),
                fixedWidthStringSize,
                type));
        }

        if (includeNulls) {
            testData.pop_back();
            testData.push_back(std::nullopt);
        }
        if (sorted) {
            testData.sort();
        }
        return testData;
    }

    template<typename T>
    VectorGeneratedData<T> genTestDataWithSequences(
        int32_t count,
        int32_t cardinality,
        bool isSorted,
        bool includeNulls,
        int32_t sequenceCount,
        int32_t sequenceLength,
        bool useFullTypeRange = true,
        const std::optional<uint32_t> &seed = std::nullopt,
        const std::optional<uint32_t> &fixedWidthStringSize = std::nullopt) {
        const bool useSequences = sequenceCount > 0 && sequenceLength > 1;
        auto testData = genTestData<T>(
            cardinality,
            CppToType<T>::create(),
            includeNulls,
            isSorted,
            useFullTypeRange,
            seed,
            fixedWidthStringSize);

        VectorGeneratedData<T> outData;
        int32_t idx = 0;

        do {
            auto val =
                    getVal(testData.data(), count, idx, sequenceCount, sequenceLength);

            // Appends a sequence.
            if (useSequences) {
                const int32_t length =
                        idx + sequenceLength > count ? count - idx : sequenceLength;
                for (auto i = 0; i < length; ++i) {
                    outData.push_back(val);
                }
                idx += sequenceLength;
            } else {
                outData.push_back(val);
                ++idx;
            }
        } while (idx < count);
        return outData;
    }

    template<typename T>
    void assertVector(
        const ExpectedData<T> &expected,
        const SimpleVectorPtr<T> &outVector,
        bool dbgPrintVec = false) {
        EXPECT_EQ(expected.size(), outVector->size());

        for (vector_size_t i = 0; i < expected.size(); i++) {
            auto optionalValue = expected[i];
            auto actualIsNull = outVector->is_null_at(i);
            const bool isNull = (optionalValue == std::nullopt);

            if (isNull) {
                if (dbgPrintVec) {
                    KLOG(INFO) << "[" << i << "]:" << "NULL";
                }
                EXPECT_EQ(isNull, actualIsNull);
            } else {
                // for all-null DictionaryVector case,
                // outVector->rawValues_ could be NULL
                if (dbgPrintVec) {
                    KLOG(INFO) << "[" << i << "]:" << *optionalValue;
                }
                if constexpr (std::is_floating_point<T>::value) {
                    if (auto isNan = std::isnan(*optionalValue)) {
                        EXPECT_EQ(isNan, std::isnan(outVector->value_at(i)));
                    } else {
                        EXPECT_EQ(*optionalValue, outVector->value_at(i));
                    }
                } else {
                    EXPECT_EQ(*optionalValue, outVector->value_at(i));
                }
            }
        }
    }

    template<typename T>
    void assertExtraMetadata(
        const VectorMakerStats<T> &expectedStats,
        const SimpleVectorPtr<T> &outVector) {
        EXPECT_EQ(expectedStats.isSorted, outVector->isSorted().value());
        EXPECT_EQ(expectedStats.min.has_value(), outVector->getMin().has_value());
        EXPECT_EQ(expectedStats.max.has_value(), outVector->getMax().has_value());
        if (expectedStats.min.has_value() && outVector->getMin().has_value()) {
            EXPECT_EQ(expectedStats.min.value(), outVector->getMin().value());
        }
        if (expectedStats.max.has_value() && outVector->getMax().has_value()) {
            EXPECT_EQ(expectedStats.max.value(), outVector->getMax().value());
        }
    }

    template<typename T>
    void assertVectorAndProperties(
        const ExpectedData<T> &expected,
        const SimpleVectorPtr<T> &outVector) {
        auto expectedStats = genVectorMakerStats(expected);

        EXPECT_EQ(expected.size(), outVector->size());
        EXPECT_EQ(
            expectedStats.distinctCount(),
            outVector->get_distinct_value_count().value());
        EXPECT_EQ(expectedStats.nullCount, outVector->get_null_count().value());
        // TODO T70862959 check getStorageBytes(), getRepresentedBytes()

        assertExtraMetadata(expectedStats, outVector);
        assertVector<T>(expected, outVector);
    }

    template<typename T>
    SimpleVectorPtr<T> createAndAssert(
        const ExpectedData<T> &expected,
        VectorEncoding::Simple encoding) {
        auto pool = memory::memoryManager()->addLeafPool();
        VectorMaker maker(pool.get());

        auto vector = maker.encoded_vector(encoding, expected);
        assertVectorAndProperties(expected, vector);
        return vector;
    }

    BufferPtr make_nulls(
        vector_size_t size,
        memory::MemoryPool *pool,
        std::function<bool(vector_size_t /*row*/)> is_null_at);

    // Create a flat vector of default values of `kind` with data-dependent flags
    // being set. The last row is set to NULL.
    template<TypeKind kind>
    VectorPtr makeFlatVectorWithFlags(
        vector_size_t size,
        memory::MemoryPool *pool) {
        using T = typename TypeTraits<kind>::NativeType;
        T value;

        auto nulls = make_nulls(size, pool, [&](auto row) { return row == size - 1; });
        auto values = AlignedBuffer::allocate<T>(size, pool, value);
        auto vector = std::make_shared<FlatVector<T> >(
            pool,
            CppToType<T>::create(),
            nulls,
            size,
            std::move(values),
            std::vector<BufferPtr>(),
            /*stats*/ SimpleVectorStats<T>{value, value},
            /*distinctValueCount*/ 1,
            /*nullCount*/ 1,
            /*isSorted*/ true,
            /*represented_bytes*/ 0,
            /*storageByteCount*/ 0);
        vector->computeAndSetIsAscii(SelectivityVector(size - 1));
        return vector;
    }

    // Create a constant vector of default values of `kind` with data-dependent
    // flags being set.
    template<TypeKind kind>
    VectorPtr makeConstantVectorWithFlags(
        vector_size_t size,
        memory::MemoryPool *pool) {
        using T = typename TypeTraits<kind>::NativeType;
        T value;

        auto vector = std::make_shared<ConstantVector<T> >(
            pool,
            size,
            false,
            CppToType<T>::create(),
            std::move(value),
            /*stats*/ SimpleVectorStats<T>{value, value},
            /*represented_bytes*/ 0,
            /*storageByteCount*/ 0);
        vector->computeAndSetIsAscii(SelectivityVector(size));
        return vector;
    }

    // Create a dictionary vector of default values of `kind` with data-dependent
    // flags being set. The indices are identity mapping and the last row is set to
    // null.
    template<TypeKind kind>
    VectorPtr makeDictionaryVectorWithFlags(
        vector_size_t size,
        memory::MemoryPool *pool) {
        using T = typename TypeTraits<kind>::NativeType;
        using TVariant = typename pollux::detail::VariantTypeTraits<kind>::native_type;
        T value;
        TVariant variant{value};

        auto nulls = make_nulls(size, pool, [&](auto row) { return row == size - 1; });
        auto base =
                BaseVector::create_constant(CppToType<T>::create(), variant, size, pool);
        auto vector = std::make_shared<DictionaryVector<T> >(
            pool,
            nulls,
            size,
            base,
            make_indices(
                size, [](auto row) { return row; }, pool),
            /*stats*/ SimpleVectorStats<T>{value, value},
            /*distinctValueCount*/ 1,
            /*nullCount*/ 1,
            /*isSorted*/ true,
            /*represented_bytes*/ 0,
            /*storageByteCount*/ 0);
        vector->computeAndSetIsAscii(SelectivityVector(size - 1));
        return vector;
    }

    // Create a flat map vector of default values of `keyKind` and `valueKind` with
    // data-dependent flags being set. The last row of the map vector and of its
    // keys and values vectors are set to null.
    template<TypeKind keyKind, TypeKind valueKind>
    VectorPtr makeMapVectorWithFlags(vector_size_t size, memory::MemoryPool *pool) {
        using K = typename TypeTraits<keyKind>::NativeType;
        using V = typename TypeTraits<valueKind>::NativeType;
        K key;
        V value;

        auto keys = makeFlatVectorWithFlags<keyKind>(size, pool);
        auto values = makeDictionaryVectorWithFlags<valueKind>(size, pool);

        auto offsets = allocateOffsets(size, pool);
        auto *rawOffsets = offsets->asMutable<vector_size_t>();
        auto sizes = allocateSizes(size, pool);
        auto *rawSizes = sizes->asMutable<vector_size_t>();
        for (auto i = 0; i < size; ++i) {
            rawOffsets[i] = i;
            rawSizes[i] = 1;
        }

        auto nulls = make_nulls(size, pool, [&](auto row) { return row == size - 1; });
        auto vector = std::make_shared<MapVector>(
            pool,
            MAP(CppToType<K>::create(), CppToType<V>::create()),
            nulls,
            size,
            offsets,
            sizes,
            keys,
            values,
            /*nullCount*/ 1,
            /*sortedKeys*/ true);
        return vector;
    }

    // Create a vector through createVector and verify that data-dependent flags
    // are set. Then call ensure_writable() or prepare_for_reuse() through makeMutable
    // and verify that data-dependent flags are cleared after the call. mutatedRows
    // specifies the rows that are mutated by makeMutable where the asciiness flag
    // should be cleared (if exists). If createVector creates an array or map
    // vector, mutatedRows should select into its element vector.
    void checkVectorFlagsReset(
        const std::function<VectorPtr()> &createVector,
        const std::function<void(VectorPtr &)> &makeMutable,
        const SelectivityVector &mutatedRows);
} // namespace kumo::pollux::test
